Greptile Summary

This PR pins all five TE-JAX Expert Parallelism FFI helpers (ep_prepare, ep_dispatch_fwd/bwd, ep_combine_fwd/bwd) to XLA's high-priority collective stream via a new _on_collective_stream decorator, preventing the XLA scheduler from interleaving them with native NCCL collectives. A version gate (is_collective_stream_supported, JAX >= 0.10.0) is added to version_utils.py so older JAX versions see no change.

• transformer_engine/jax/cpp_extensions/ep.py — introduces _on_collective_stream, applied at module-load time; lazily imports jax.experimental.compute_on.compute_on only when supported.
• transformer_engine/jax/version_utils.py — adds _COLLECTIVE_STREAM_MIN_JAX_VERSION = \"0.10.0\" and is_collective_stream_supported(), exported in __all__.
• tests/jax/test_multi_process_ep.py — adds a forward-only HLO test that compiles a dispatch/combine trace and asserts every te_ep_* custom call carries _xla_stream_annotation=\"collective\"; test is skipped on unsupported JAX.

Confidence Score: 5/5

Safe to merge — the feature is gated behind a version check so older JAX installations see no behavioural change, and the new test validates the annotation on the forward path.

The change is an opt-in scheduling hint applied at module import time. The version guard correctly isolates older runtimes, the decorator mechanism is straightforward, and the new HLO test confirms the annotation appears in compiled output. No logic errors or API misuses were found.

The ep_dispatch_bwd and ep_combine_bwd backward ops are decorated but their annotation is not verified by any test — worth adding a backward-path assertion to test_z_no_unexpected_reshard_in_hlo_bwd or a companion test.

Important Files Changed

Sequence Diagram

sequenceDiagram
    participant TE as TE ep.py import
    participant VU as version_utils
    participant XLA as XLA scheduler

    TE->>VU: "is_collective_stream_supported()?"
    VU-->>TE: "JAX >= 0.10.0 → True"
    TE->>TE: "compute_on(gpu_stream:collective)(func) wraps all 5 EP helpers"

    Note over TE,XLA: At JIT compile time

    TE->>XLA: "EpDispatchPrimitive.bind() with _xla_stream_annotation=collective"
    XLA-->>XLA: "schedule on NCCL collective stream"
    TE->>XLA: "EpCombinePrimitive.bind() with _xla_stream_annotation=collective"
    XLA-->>XLA: "serialised with native all-reduce / all-to-all ops"

_{Reviews (5): Last reviewed commit: "Merge branch 'main' into ep/jax_col_stre..." | Re-trigger Greptile}

/te-ci JAX L1

/te-ci JAX L1